import cv2 as cv
# MediaPipe是一个用于构建机器学习管道的框架，用于处理视频、音频等时间序列数据。这个跨平台框架适用于桌面/服务器、Android、iOS和嵌入式设备，如Raspberry Pi和Jetson Nano。
import mediapipe as mp
import time
import math

class poseDetector():
    def __init__(self,static_image_mode=False, # 静态图模式，False代表置信度高时继续跟踪，True代表实时跟踪检测新的结果
                    upper_body_only=False,  # 是否只检测上半身
                    model_complexity=1,
                    smooth_landmarks=True,  # 平滑，一般为True
                    min_detection_confidence=0.5, # 检测置信度# 检测置信度大于0.5代表检测到了，若此时跟踪置信度大于0.5就继续跟踪，小于就沿用上一次，避免一次又一次重复使用模型
                    min_tracking_confidence=0.5):  # 跟踪置信度
        self.static_image_mode = static_image_mode
        self.upper_body_only = upper_body_only
        self.model_complexity=model_complexity
        self.smooth_landmarks = smooth_landmarks
        self.min_detection_confidence = min_detection_confidence
        self.min_tracking_confidence = min_tracking_confidence

        # mediapipe.solutions.drawing_utils.draw_landmarks()绘制关键点的连线
        self.mpDraw = mp.solutions.drawing_utils
        # mediapipe.solutions.pose.Pose()姿态关键点检测函数
        self.mpPose = mp.solutions.pose
        self.pose = self.mpPose.Pose(self.static_image_mode,self.upper_body_only,self.model_complexity,self.smooth_landmarks,self.min_detection_confidence,self.min_tracking_confidence)

    # 检测关键点方法    
    def findPose(self, frame, draw=True):
        self.frame_RGB = cv.cvtColor(frame, cv.COLOR_BGR2RGB)
        # 3.将图像传给姿态识别模型
        self.res = self.pose.process(self.frame_RGB)
        if self.res.pose_landmarks:
            if draw:
                # 绘制姿态坐标点，img为画板，传入姿态点坐标，坐标连线
                # mediapipe.solutions.drawing_utils.draw_landmarks()绘制手部关键点的连线
                # mpDraw.draw_landmarks(frame,res.pose_landmarks)
                self.mpDraw.draw_landmarks(frame, self.res.pose_landmarks, self.mpPose.POSE_CONNECTIONS)
            return frame

        
    #关键点信息，是否绘制关键点连线    
    def findPosition(self, frame, draw=True):
        self.lmList = []
        if self.res.pose_landmarks:
            for id,lm in enumerate(self.res.pose_landmarks.landmark):
                h, w, c = frame.shape
                cx, cy = int(lm.x * w), int(lm.y * h)
                self.lmList.append([id,cx,cy])
                if draw:
                    cv.circle(frame, (cx, cy), 10, (255, 0, 0), cv.FILLED)
        return self.lmList
    
    #计算运动时关键点角度    
    def findAngle(self,frame,p1,p2,p3,draw=True):
            # 关键点坐标
            x1,y1 = self.lmList[p1][1:]
            x2,y2 = self.lmList[p2][1:]
            x3,y3 = self.lmList[p3][1:]

            # 关键点角度 atan2 方法返回一个 -pi 到 pi 之间的数值，表示点 (x, y) 对应的偏移角度。这是一个逆时针角度，以弧度为单位，正X轴和点 (x, y) 与原点连线 之间
            angle = math.degrees(math.atan2(y3-y2,x3-x2)-math.atan2(y1-y2,x1-x2))

            # 绘制关键点
            if draw:
                cv.line(frame,(x1,y1),(x2,y2),(255,0,255),3)
                cv.line(frame,(x3,y3),(x2,y2),(255,0,255),3)
                cv.circle(frame,(x1,y1),10,(0,0,255),cv.FILLED)
                cv.circle(frame,(x1,y1),15,(0,0,255),2)
                cv.circle(frame,(x2,y2),10,(0,0,255),cv.FILLED)
                cv.circle(frame,(x2,y2),15,(0,0,255),2)
                cv.circle(frame,(x3,y3),10,(0,0,255),cv.FILLED)
                cv.circle(frame,(x3,y3),15,(0,0,255),2)
            return angle